% TWO-DIMENSIONAL, MESH-CENTERED, MULTIGROUP DIFFUSION
% 
% TO DO:  
%   * reimplement incident boundary source and verify
%   * expand to general group structure and verify
%   * is it worth using an iterative solver?
%   * vectorize

clear; 
%clc; %clf;

disp('-------------------------------------------')
disp('--     beginning 2D diffusion solver     --')
disp('-------------------------------------------')

%clear
%---- INPUT FILE (JUST AN "M" FILE) ---------------------------------------
%twodbenchmark_input  % 2-D IAEA benchmark, very close
%simple2d_input      % simple 1 box problem for reference
%samp2d1gfixed
%test_case_by_hand
%clear, test_hand_2
%clear, test_hand_two_group
%reconstruct
%twodbenchmark_input
%twod_biblis_input  % 2-d biblis test cast
%twod_koeberg_input % 2-d 4-group koeberg input
checkerboard
tic

%---- GENERATE THE COEFFICIENTS -------------------------------------------
[AB, AL, AC, AR, AT, Si, Sh, dx, dy, dv, dc, siga, sct, chi] = ...
    mesh_center_2d_coef( dat, numg, numm, xcm, xfm, ...
       ycm, yfm, mt, src, it, [BCL BCR BCB BCT], IBSL, IBSR, IBSB, IBST );
% can always use .._coef()


%---- SOLVE----------------------------------------------------------------
N = sum(xfm); M = sum(yfm);
% pad end of AB,AL
AB = [AB' zeros(numg,N)]';
AL = [AL' zeros(numg,1)]';
% pad top of AR,AT
AT = [zeros(numg,N) AT']';
AR = [zeros(numg,1) AR']';
phi = zeros( N*M, numg );

% build cell structure of sparse matrices
for g = 1:numg
   A{g} = spdiags([AB(:,g) AL(:,g) AC(:,g) AR(:,g) AT(:,g)], ...
       [-N -1 0 1 N], N*M,N*M);
end
% phi(:,1) = A1 \ ( S(:,1) +  1/keff * nusig(:,2).*phi(:,2) );
if it==0 % fixed source, no multiplication --------------------------------
    phieps = 1e-8;  mxit = 200;  it = 0; phierr = 1; 
    while (phierr > phieps && i < mxit)
        phi0 = phi;
        phi(:,1) = A{1} \ Si(:,1);
        for g = 2:numg
            scsrc = zeros(N*M,1); % reset
            for gg = 1:g-1
                scsrc(:,1) = scsrc(:,1) + sct(:,gg,g).*phi(:,gg);
            end
            phi(:,g) = A{g} \ ( Si(:,g) + scsrc(:,1) );
        end
        phierr = norm( phi0 - phi );
        i = i + 1;
    end
elseif it==-1 % fixed source with multiplication --------------------------
    nusig = Sh;  % just rename it right away to avoid confusion
    s = 0*ones(N*M,1)/(ycm(end)+xcm(end));  % uniform guess
    phieps = 1e-10;  mxit = 200;  i = 0; phierr = 1; 
    while (phierr > phieps && i < mxit)
        phi0 = phi;
        phi(:,1) = A{1} \ ( Si(:,1) + chi(:,1).*s(:,1)/keff ) ;
        for g = 2:numg
            scsrc = zeros(N*M,1); % reset
            for gg = 1:g-1
                scsrc(:,1) = scsrc(:,1) + sct(:,gg,g).*phi(:,gg);
            end
            phi(:,g) = A{g} \ ( Si(:,g) + scsrc(:,1) + chi(:,g).*s(:,1)/keff  );
        end
        s(:,1) = 0; % reset fission source
        for g = 1:numg
            s(:,1) = s(:,1) + Sh(:,g).*phi(:,g);
        end         
        phierr = norm( phi0 - phi );
        i = i + 1;
    end
else % eigenvalue
    disp('begin keff iteration')
    %nusig = S;  % just rename it right away to avoid confusion
    s = ones(N*M,1)/(ycm(end)+xcm(end));  % uniform guess
    keff    = 1.0;    % guess the initial keff
    errK    = 1;      errS    = 1; 
    epsK    = 1e-8;   epsS = 1e-6;
    iter    = 0;      itmx = 300;
    phi     = zeros(N*M,numg); 

    while  ( ( errK>epsK ) || (errS>epsS) ) && iter < itmx
        % set to solve 1st group flux
        s = s/sum(s);
        phi(:,1) = A{1} \ (chi(:,1).*s(:,1)/keff) ;  % assumes all fiss n's fast
        if numg > 1
            for g = 2:numg
                scsrc = zeros(N*M,1); % reset
                % compute scattering source
                for gg = 1:numg % could put a variable bound for upscatter
                    scsrc(:,1) = scsrc(:,1) + sct(:,g,gg).*phi(:,gg);
                end
                phi(:,g) = A{g} \ (chi(:,g).*s(:,1)/keff+scsrc(:,1) ); 
            end   
        end
        sold = s; 
        kold = keff;   
        s(:,1) = 0; % reset fission source
        for g = 1:numg
            s(:,1) = s(:,1) + Sh(:,g).*phi(:,g);
        end 
        keff = sum(dv.*s)*kold/sum(dv.*sold);
        indx = find(s);
        errS = max( abs((s(indx)-sold(indx))./s(indx)) );
        errK = abs( (keff-kold)/keff );
        iter = iter + 1;  % number of iterations
        if ( mod(iter,1)== 0)
            disp([' iter = ',num2str(iter), ' keff = ',num2str(keff)])
            disp([' errK = ',num2str(errK), ' errS = ',num2str(errS)])
        end
    end
    disp([ 'final result: keff = ',num2str(keff)])
    disp([ '              iter = ',num2str(iter)])
    keff
end

 toc
% 
% [CRL,CRR,CRB,CRT,totLL,totLR,totLB,totLT] = ...
%     cresp_mc(phi,dx,dy,dc,mt,xfm,yfm,numg,IBSB,IBSL,IBSR,IBST);
% 
% ord=1;
% % % approximate responses,1.048320650377119
% CRLlegcoef = vec2leg(CRL,length(CRL),1,ord)';
% CRLappx    = vec2leg( CRLlegcoef, length(CRL),xfm,ord);
% CRRlegcoef = vec2leg(CRR,length(CRR),1,ord)';
% CRRappx    = vec2leg( CRRlegcoef, length(CRR),xfm,ord);
% CRBlegcoef = vec2leg(CRB,length(CRB),1,ord)';
% CRBappx    = vec2leg( CRBlegcoef, length(CRB),xfm,ord);
% CRTlegcoef = vec2leg(CRT,length(CRT),1,ord)';
% CRTappx    = vec2leg( CRTlegcoef, length(CRT),xfm,ord);
% 
% ff = [ CRLlegcoef CRRlegcoef CRBlegcoef CRTlegcoef ]'
% 
% gain = phi'*Sh;
% abs  = phi'*siga;
% leak = totLL*(1-(BCL>=1))+totLR*(1-(BCR>=1))+...
%        totLT*(1-(BCT>=1))+totLB*(1-(BCB>=1));
% 
% alt_keff = gain/(abs+leak)
% 
%max(phi(:,1)./phi(:,2))
xx(1)=0; 
for i = 2:length(dy)+1;
    xx(i)=xx(i-1)+dy(i-1);
end
x=0.5*(xx(2:end)+xx(1:end-1));
p1 = reshape(2*phi(:,1),M,N); p1x = p1(1,:); hold on, figure(1), plot(x,p1x)
FR = phi(:,1).*Sh(:,1) + phi(:,2).*Sh(:,2);
FR = FR / norm(FR);
% yy(1)=0; 
% for i = 2:length(dy)+1;
%     yy(i)=yy(i-1)+dy(i-1);
% end
% y=0.5*(yy(2:end)+yy(1:end-1));
% 
% for g = 1:numg
%     figure(g)
%     contourf(reshape(phi(:,g),length(dx),length(dy)),100)
%     shading flat
%     colorbar
%     square axis;
%     title(['\phi',num2str(g)]), xlabel('y [cm]'), ylabel('x [cm]')
% end
%phi
%phi2 =reshape(phi,9,9);
%contourf( phi2,100 ), colorbar, shading flat

